#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
启语(Chim)编译器 - 代码生成器
"""
import sys
import os
from parser_ import (
    Struct,
    Function,
    VariableDeclaration,
    MatchStatement,
    ForLoop,
    ReturnStatement,
    CallExpression,
    BinaryExpression,
    UnaryExpression,
    Identifier,
    Literal,
    ArrayLiteral,
    MapLiteral,
    ImportStatement,
    MemberAccess,
    GroupBlock,
    TuplePattern,
    RangePattern,
    SnapshotExpression,
    HandleExpression,
    ArenaCall,
)

class CodeGenerator:
    def __init__(self, ast, dialect="safe"):
        self.ast = ast
        self.dialect = dialect
        self.output = []
        self.indent_level = 0
        self.functions = {}
        self.structs = {}
        self.variables = {}
    
    def indent(self):
        return "    " * self.indent_level
    
    def emit(self, code):
        self.output.append(self.indent() + code)
    
    def generate(self):
        """生成代码"""
        self.visit_program(self.ast)
        return "\n".join(self.output)
    
    def visit_program(self, node):
        """访问程序节点"""
        # 添加文件头部注释
        self.emit("""# Generated by Chim Compiler
# 启语(Chim)语言
""")
        
        if self.dialect == "system":
            self.emit("import sys")
            self.emit("import os")
        else:
            self.emit("import sys")
            self.emit("import os")
        self.emit("sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'runtime_py'))")
        self.emit("sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'compiler'))")
        self.emit("from chim_runtime import 输出, 范围, send, recv")
        # 将导入语句从AST前置输出
        for stmt in node.statements:
            if isinstance(stmt, ImportStatement):
                py_mod = 'chym.' + stmt.module_path.replace("::", ".")
                if stmt.alias:
                    self.emit(f"import {py_mod} as {stmt.alias}")
                else:
                    self.emit(f"import {py_mod}")
        self.emit("")
        
        # 先收集所有函数和结构体定义
        for stmt in node.statements:
            if isinstance(stmt, Struct):
                self.structs[stmt.name] = stmt
            elif isinstance(stmt, Function):
                self.functions[stmt.name] = stmt
        
        # 生成结构体定义
        for struct_name, struct_node in self.structs.items():
            self.visit_struct(struct_node)
            self.emit("")
        
        # 生成函数定义
        for func_name, func_node in self.functions.items():
            self.visit_function(func_node)
            self.emit("")
        
        # 添加main函数调用
        self.emit("if __name__ == '__main__':")
        self.indent_level += 1
        if 'main' in self.functions:
            self.emit("main()")
        elif '主' in self.functions:
            self.emit("主()")
        self.indent_level -= 1
    
    def visit_struct(self, node):
        """访问结构体节点，转换为Python类"""
        self.emit(f"class {node.name}:")
        self.indent_level += 1
        
        # 生成__init__方法
        self.emit("def __init__(self):")
        self.indent_level += 1
        for field_name, field_type in node.fields:
            self.emit(f"self.{field_name} = None  # 类型: {field_type}")
        self.indent_level -= 1
        self.emit("")
        
        # 生成__repr__方法
        self.emit("def __repr__(self):")
        self.indent_level += 1
        fields_repr = ", ".join(f"{field}={{self.{field}}}" for field, _ in node.fields)
        self.emit(f"return f'{node.name}({fields_repr})'")
        self.indent_level -= 1
        
        self.indent_level -= 1
    
    def visit_function(self, node):
        """访问函数节点，转换为Python函数"""
        # 处理参数
        params = [f"{param_name}" for param_name, _ in node.params]
        
        # 生成函数定义
        self.emit(f"def {node.name}({', '.join(params)}):")
        self.indent_level += 1
        
        # 添加注释
        if node.return_type:
            self.emit(f"# 返回类型: {node.return_type}")
        
        # 生成函数体
        self.visit_statements(node.body)
        
        self.indent_level -= 1
    
    def visit_statements(self, statements):
        """访问语句列表"""
        for stmt in statements:
            self.visit_statement(stmt)
    
    def visit_statement(self, stmt):
        """根据语句类型分派访问"""
        if isinstance(stmt, VariableDeclaration):
            self.visit_variable_declaration(stmt)
        elif isinstance(stmt, MatchStatement):
            self.visit_match_statement(stmt)
        elif isinstance(stmt, ForLoop):
            self.visit_for_loop(stmt)
        elif isinstance(stmt, ReturnStatement):
            self.visit_return_statement(stmt)
        elif isinstance(stmt, CallExpression):
            self.visit_call_expression(stmt)
        elif isinstance(stmt, BinaryExpression):
            # 作为语句的表达式需要处理，但一般函数调用更多
            pass
        elif stmt.__class__.__name__ == 'BreakStatement':
            self.emit("break")
        elif stmt.__class__.__name__ == 'ContinueStatement':
            self.emit("continue")
        elif isinstance(stmt, GroupBlock):
            self.visit_group_block(stmt)
    
    def visit_variable_declaration(self, node):
        """访问变量声明节点"""
        # Python没有显式类型，添加类型注释
        if node.type:
            var_def = f"{node.name} = {self.visit_expression(node.value)}  # 类型: {node.type}"
        else:
            var_def = f"{node.name} = {self.visit_expression(node.value)}"
        
        # 对于可变变量，Python没有区别，但我们记录它
        if node.is_mutable:
            self.variables[node.name] = "mutable"
        else:
            self.variables[node.name] = "immutable"
        
        self.emit(var_def)
    
    def visit_match_statement(self, node):
        """访问匹配语句节点，支持Swift风格增强 - 元组、范围、guard"""
        # 为表达式赋值给临时变量
        temp_var = "_match_expr"
        self.emit(f"{temp_var} = {self.visit_expression(node.expression)}")
        
        has_default = False
        first_case = True
        
        for item in node.cases:
            # 解包: 新格式包含guard
            if len(item) == 3:
                pattern, guard, body = item
            else:
                # 向后兼容旧格式
                pattern, body = item
                guard = None
            
            # 检查是否是默认分支
            is_default = isinstance(pattern, Literal) and pattern.value_type == "DEFAULT"
            is_underscore_default = isinstance(pattern, Identifier) and pattern.name == "_"
            
            if is_default or is_underscore_default:
                has_default = True
                # guard在默认分支不适用
                self.emit("else:")
            else:
                # 生成条件
                condition = self.generate_match_condition(temp_var, pattern)
                
                # 添加guard条件
                if guard:
                    guard_str = self.visit_expression(guard)
                    condition = f"({condition}) and ({guard_str})"
                
                if first_case:
                    self.emit(f"if {condition}:")
                    first_case = False
                else:
                    self.emit(f"elif {condition}:")
            
            # 生成case体
            self.indent_level += 1
            self.visit_statements(body)
            self.indent_level -= 1
        
        # 确保有默认分支(安全方言要求)
        if not has_default and self.dialect == "safe":
            self.emit("else:")
            self.indent_level += 1
            self.emit("pass  # 安全方言要求默认分支")
            self.indent_level -= 1
    
    def generate_match_condition(self, temp_var, pattern):
        """生成匹配条件 - 支持元组、范围等模式"""
        # 元组模式: (a, b, c)
        if isinstance(pattern, TuplePattern):
            conditions = []
            bindings = []  # 保存需要绑定的变量
            for i, elem in enumerate(pattern.elements):
                # 如果是标识符，不是字面量，则是模式变量(任意匹配)
                if isinstance(elem, Identifier):
                    # 模式变量 - 同时赋值
                    var_name = elem.name
                    bindings.append(f"{var_name} := {temp_var}[{i}]")
                else:
                    # 字面值匹配
                    value_str = self.visit_expression(elem)
                    conditions.append(f"{temp_var}[{i}] == {value_str}")
            
            # 先进行绑定，再检查条件
            if bindings:
                # Python 3.8+ walrus operator
                for binding in bindings:
                    conditions.insert(0, f"({binding}) or True")
            
            return " and ".join(conditions) if conditions else "True"
        
        # 范围模式: a..b 或 a...b
        elif isinstance(pattern, RangePattern):
            start_str = self.visit_expression(pattern.start)
            end_str = self.visit_expression(pattern.end)
            if pattern.inclusive:
                # 闭区间 [start, end]
                return f"({start_str} <= {temp_var} <= {end_str})"
            else:
                # 半开区间 [start, end)
                return f"({start_str} <= {temp_var} < {end_str})"
        
        # 模式变量绑定: 案例 x ->
        elif isinstance(pattern, Identifier) and pattern.name not in ('_', '默认', 'true', 'false', '真', '假'):
            # 这是一个模式变量，需要绑定
            var_name = pattern.name
            # 使用walrus operator绑定并总是匹配
            return f"({var_name} := {temp_var}) or True"
        
        # 普通值匹配
        else:
            pattern_str = self.visit_expression(pattern)
            return f"{temp_var} == {pattern_str}"
    
    def visit_for_loop(self, node):
        """访问for循环节点，转换为Python的for循环"""
        iterable_str = self.visit_expression(node.iterable)
        self.emit(f"for {node.variable} in {iterable_str}:")
        self.indent_level += 1
        self.visit_statements(node.body)
        self.indent_level -= 1
    
    def visit_return_statement(self, node):
        """访问返回语句节点"""
        if node.expression:
            expr_str = self.visit_expression(node.expression)
            self.emit(f"return {expr_str}")
        else:
            self.emit("return")
    
    def visit_expression(self, expr):
        """根据表达式类型分派访问"""
        if isinstance(expr, CallExpression):
            return self.visit_call_expression_as_expr(expr)
        elif isinstance(expr, BinaryExpression):
            return self.visit_binary_expression_as_expr(expr)
        elif isinstance(expr, UnaryExpression):
            return self.visit_unary_expression_as_expr(expr)
        elif isinstance(expr, Identifier):
            return expr.name
        elif isinstance(expr, MemberAccess):
            return f"{self.visit_expression(expr.object)}.{expr.member.name}"
        elif isinstance(expr, Literal):
            return self.visit_literal_as_expr(expr)
        elif isinstance(expr, ArrayLiteral):
            return self.visit_array_literal_as_expr(expr)
        elif isinstance(expr, MapLiteral):
            return self.visit_map_literal_as_expr(expr)
        elif isinstance(expr, SnapshotExpression):
            return self.visit_snapshot_expression(expr)
        elif isinstance(expr, HandleExpression):
            return self.visit_handle_expression(expr)
        elif isinstance(expr, ArenaCall):
            return "None  # Arena分配"
        elif isinstance(expr, TuplePattern):
            # 元组模式作为表达式(在匹配中不应该直接转换)
            elements = [self.visit_expression(e) for e in expr.elements]
            return f"({', '.join(elements)})"
        elif isinstance(expr, RangePattern):
            # 范围模式不应直接作为表达式，但如果需要...
            return f"range({self.visit_expression(expr.start)}, {self.visit_expression(expr.end)})"
        else:
            raise ValueError(f"未知的表达式类型: {type(expr)}")
    
    def visit_group_block(self, node):
        """访问组块 - 生成Python上下文管理器"""
        group_name = node.name.name if isinstance(node.name, Identifier) else str(node.name)
        
        # 生成注释
        self.emit(f"# 组: {group_name} (批量释放)")
        self.emit("try:")
        self.indent_level += 1
        
        # 生成组内语句
        self.visit_statements(node.body)
        
        self.indent_level -= 1
        self.emit("finally:")
        self.indent_level += 1
        self.emit(f"pass  # 组 '{group_name}' 结束，自动释放资源")
        self.indent_level -= 1
    
    def visit_snapshot_expression(self, node):
        """访问快照表达式 - 生成只读副本"""
        target = self.visit_expression(node.target)
        # 在Python中，使用copy.deepcopy创建快照
        return f"list({target}) if isinstance({target}, list) else {target}  # 快照"
    
    def visit_handle_expression(self, node):
        """访问句柄表达式 - 生成引用"""
        target = self.visit_expression(node.target)
        # 在Python中，直接引用
        return f"{target}  # 句柄"
    
    def visit_call_expression(self, node):
        """访问函数调用节点（作为语句）"""
        self.emit(self.visit_call_expression_as_expr(node))
    
    def visit_call_expression_as_expr(self, node):
        """访问函数调用节点（作为表达式）"""
        # 特例：io.输出 → print
        if isinstance(node.function, MemberAccess):
            if isinstance(node.function.object, Identifier) and node.function.object.name == 'io' and node.function.member.name == '输出':
                args = [str(self.visit_expression(arg)) for arg in node.arguments]
                return f"print({', '.join(args)})"
        func_name = self.visit_expression(node.function)
        args = [str(self.visit_expression(arg)) for arg in node.arguments]
        return f"{func_name}({', '.join(args)})"
    
    def visit_binary_expression_as_expr(self, node):
        """访问二元表达式节点（作为表达式）"""
        left_str = self.visit_expression(node.left)
        right_str = self.visit_expression(node.right)
        
        # 处理二元运算符映射
        op_map = {
            "+" : "+",
            "-" : "-",
            "*" : "*",
            "/" : "/",
            "%" : "%",
            "^" : "**",
            "==" : "==",
            "!=" : "!=",
            ">" : ">",
            "<" : "<",
            ">=" : ">=",
            "<=" : "<=",
            "&&" : " and ",
            "||" : " or "
        }
        
        operator = op_map.get(node.operator, node.operator)
        if node.operator == "<-":
            return f"send({left_str}, {right_str})"
        
        # 对于布尔运算符，需要特殊处理括号
        if node.operator in ["&&", "||"]:
            return f"({left_str}){operator}({right_str})"
        else:
            return f"{left_str} {operator} {right_str}"
    
    def visit_unary_expression_as_expr(self, node):
        """访问一元表达式节点（作为表达式）"""
        operand_str = self.visit_expression(node.operand)
        
        # 处理一元运算符映射
        if node.operator == "!":
            return f"not {operand_str}"
        elif node.operator == "-":
            return f"-{operand_str}"
        elif node.operator == "<-":
            return f"recv({operand_str})"
        
        raise ValueError(f"未知的一元运算符: {node.operator}")
    
    def visit_literal_as_expr(self, node):
        """访问字面量节点（作为表达式）"""
        if node.value_type == "INTEGER":
            return node.value
        elif node.value_type == "FLOAT":
            return node.value
        elif node.value_type == "STRING":
            return repr(node.value)
        elif node.value_type == "CHAR":
            return repr(node.value)
        elif node.value_type == "BOOLEAN":
            return "True" if node.value else "False"
        elif node.value_type == "NIL":
            return "None"
        elif node.value_type == "DEFAULT":
            return "_"

        raise ValueError(f"未知的字面量类型: {node.value_type}")
    
    def visit_array_literal_as_expr(self, node):
        """访问数组字面量节点（作为表达式）"""
        elements = [str(self.visit_expression(elem)) for elem in node.elements]
        return f"[{', '.join(elements)}]"
    
    def visit_map_literal_as_expr(self, node):
        """访问映射字面量节点（作为表达式）"""
        entries = []
        for key, value in node.entries:
            key_str = self.visit_expression(key)
            value_str = self.visit_expression(value)
            entries.append(f"{key_str}: {value_str}")
        return f"{{{', '.join(entries)}}}"

# 导出函数
def generate_code(ast, dialect="safe"):
    """生成代码的主函数"""
    generator = CodeGenerator(ast, dialect)
    return generator.generate()

def generate_code_from_ir(function_names):
    lines = []
    lines.append("import sys")
    lines.append("import os")
    lines.append("sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'runtime_py'))")
    lines.append("sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'compiler'))")
    lines.append("from chim_runtime import send, recv")
    for name in function_names:
        lines.append(f"def {name}():")
        lines.append("    pass")
    return "\n".join(lines)

def generate_code_from_ir_full(mod):
    lines = []
    lines.append("import sys")
    lines.append("import os")
    lines.append("sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'runtime_py'))")
    lines.append("sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'compiler'))")
    lines.append("from chim_runtime import send, recv")
    # 结构：mod.funcs[i]，每个函数包含 name, args, arg_count, return_type
    for i in range(mod.func_count):
        f = mod.funcs[i]
        name = f.name.decode('utf-8') if f.name else 'fn'
        params = []
        if f.arg_count > 0 and f.args:
            for j in range(f.arg_count):
                arg = f.args[j]
                arg_name = arg.name.decode('utf-8') if arg.name else f"arg{j}"
                params.append(arg_name)
        lines.append(f"def {name}({', '.join(params)}):")
        lines.append("    pass")
    return "\n".join(lines)
